import zmq
import numpy as np
import cv2
import json
from multiprocessing import shared_memory
import subprocess
import atexit
import os
import sys
import time
import uuid
import platform
import traceback

class ImageProcessor:
    def __init__(self):
        self.rust_process = None
        self.ctx = zmq.Context()
        self.socket = self.ctx.socket(zmq.REQ)
        self.socket.RCVTIMEO = 10000  # 10秒超时
        
        # 确定IPC路径（跨平台）
        self.ipc_path = self._get_ipc_path()
        self._start_rust_process()
        atexit.register(self.cleanup)

    def _get_ipc_path(self):
        """获取适合当前操作系统的IPC路径"""
        if platform.system() == "Windows":
            return "tcp://127.0.0.1:5555"
        else:
            return "ipc:///tmp/image_processor.ipc"
    
    def _start_rust_process(self):
        """启动Rust处理进程"""
        # 获取Rust可执行文件路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        rust_dir = os.path.join(current_dir, "..", "rust_processor")
        target_dir = "debug"  # 或 "release" 当使用--release构建时
        
        # 构建可执行文件名（跨平台）
        executable = "rust_processor"
        if platform.system() == "Windows":
            executable += ".exe"
        
        rust_bin = os.path.join(rust_dir, "target", target_dir, executable)
        
        # 如果未构建，则尝试构建
        if not os.path.exists(rust_bin):
            print("Building Rust processor...")
            build_result = subprocess.run(
                ["cargo", "build"],
                cwd=rust_dir,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            if build_result.returncode != 0:
                raise RuntimeError(
                    f"Failed to build Rust processor:\n"
                    f"STDOUT: {build_result.stdout.decode()}\n"
                    f"STDERR: {build_result.stderr.decode()}"
                )
        
        # 启动Rust进程
        self.rust_process = subprocess.Popen(
            [rust_bin],
            cwd=rust_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        # 等待进程启动
        time.sleep(2)
        if self.rust_process.poll() is not None:
            # 进程已退出，获取错误信息
            stdout, stderr = self.rust_process.communicate()
            raise RuntimeError(
                f"Rust processor failed to start:\nSTDOUT: {stdout}\nSTDERR: {stderr}"
            )
        
        # 连接到IPC
        try:
            self.socket.connect(self.ipc_path)
            print(f"Connected to Rust processor at {self.ipc_path}")
        except zmq.ZMQError as e:
            raise RuntimeError(f"Failed to connect to IPC: {str(e)}")

    def process(self, command: str, image: np.ndarray, params: dict = None) -> np.ndarray:
        """处理图像并返回结果"""
        # 生成唯一共享内存ID
        shmem_id = f"pyimg_{os.getpid()}_{uuid.uuid4().hex}"
        
        # 准备图像数据
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)
        
        # 编码为PNG
        success, buffer = cv2.imencode(".png", image)
        if not success:
            raise ValueError("Failed to encode image as PNG")
        
        buffer_bytes = buffer.tobytes()
        
        # 创建共享内存
        try:
            shm = shared_memory.SharedMemory(
                name=shmem_id, 
                create=True, 
                size=len(buffer_bytes)
            )
            # 复制数据到共享内存
            shm.buf[:len(buffer_bytes)] = buffer_bytes
        except Exception as e:
            raise RuntimeError(f"Shared memory creation failed: {str(e)}")

        # 发送控制消息
        ctrl_msg = {
            "command": command,
            "shmem_id": shmem_id,
            "width": image.shape[1],
            "height": image.shape[0],
            "params": params
        }
        
        try:
            self.socket.send_string(json.dumps(ctrl_msg))
        except zmq.ZMQError as e:
            self._cleanup_shmem(shm)
            raise RuntimeError(f"Failed to send control message: {str(e)}")

        # 接收响应
        try:
            response = self.socket.recv_string()
        except zmq.Again:
            self._cleanup_shmem(shm)
            raise TimeoutError("Processing timeout after 10 seconds")
        except Exception as e:
            self._cleanup_shmem(shm)
            raise RuntimeError(f"Communication error: {str(e)}")

        resp_data = json.loads(response)
        if resp_data["status"] != "success":
            self._cleanup_shmem(shm)
            error_msg = resp_data.get("error", "Unknown error")
            raise RuntimeError(f"Processing error: {error_msg}")

        # 获取结果
        result_shm_id = resp_data["result_shmem_id"]
        try:
            result_shm = shared_memory.SharedMemory(name=result_shm_id)
            # 获取整个共享内存内容
            result_data = bytes(result_shm.buf)
            result_shm.close()
            result_shm.unlink()
        except Exception as e:
            self._cleanup_shmem(shm)
            raise RuntimeError(f"Result shmem access failed: {str(e)}")

        # 清理输入共享内存
        self._cleanup_shmem(shm)

        # 解码图像
        result_img = cv2.imdecode(
            np.frombuffer(result_data, dtype=np.uint8), 
            cv2.IMREAD_UNCHANGED
        )
        
        if result_img is None:
            raise RuntimeError("Failed to decode result image")
            
        return result_img

    def _cleanup_shmem(self, shm: shared_memory.SharedMemory):
        """安全清理共享内存"""
        try:
            shm.close()
            shm.unlink()
        except Exception as e:
            print(f"Warning: Failed to cleanup shared memory: {str(e)}")

    def cleanup(self):
        """清理资源"""
        try:
            self.socket.close()
            self.ctx.term()
        except:
            pass
            
        if self.rust_process and self.rust_process.poll() is None:
            try:
                self.rust_process.terminate()
                self.rust_process.wait(timeout=2)
            except:
                try:
                    self.rust_process.kill()
                except:
                    pass

if __name__ == "__main__":
    # 测试代码
    processor = ImageProcessor()
    
    # 创建测试图像
    test_img = np.zeros((480, 640, 3), dtype=np.uint8)
    cv2.putText(test_img, "Test Image", (50, 240), 
                cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3)
    
    try:
        # 测试灰度处理
        print("Processing grayscale...")
        gray = processor.process("grayscale", test_img)
        cv2.imwrite("result_grayscale.jpg", gray)
        
        # 测试旋转
        print("Processing rotate90...")
        rotated = processor.process("rotate90", test_img)
        cv2.imwrite("result_rotated.jpg", rotated)
        
        # 测试模糊
        print("Processing blur...")
        blurred = processor.process("blur", test_img)
        cv2.imwrite("result_blurred.jpg", blurred)
        
        # 测试边缘检测
        print("Processing canny_edge...")
        edges = processor.process("canny_edge", test_img)
        cv2.imwrite("result_edges.jpg", edges)
        
        print("Processing completed. Check result images.")
    except Exception as e:
        print(f"Error during processing: {str(e)}")
        traceback.print_exc()